# coding=utf-8
# @author:      ChengJing
# @name:        ICNN.py
# @datetime:    2022/2/3 23:17
# @software:    PyCharm
# @description:


import torch
import torch.nn as nn
from model.improve_location.ImproveFeatures import LAQ


class ICNN(nn.Module):
    """
    定义简单的卷积神经网络模型进行爆管预警和定位识别
    """
    def __init__(self, g, sensors, in_channel, hidden_channel1, hidden_channel2, out_channel, hidden_feature, out_features):
        """
        Args:
            g: 流量监测点和压力监测点之间的连接关系，shape：N*M
            sensors: 压力监测点的数量
            in_channel: 输入数据的通道数
            hidden_channel1: 隐含层通道数
            hidden_channel2: 隐含层通道数
            out_channel: 卷积输出通道数
            hidden_feature: 隐含层特征
            out_features: 输出特征
        """
        super(ICNN, self).__init__()
        self.laq = LAQ(g, sensors, 3*sensors)
        self.cnn = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=hidden_channel1, kernel_size=3, padding=2),
            nn.MaxPool2d(kernel_size=3),
            nn.Conv2d(in_channels=hidden_channel1, out_channels=hidden_channel2, kernel_size=(5, 3), padding=2),
            nn.MaxPool2d(kernel_size=(5, 3)),
            nn.Conv2d(in_channels=hidden_channel2, out_channels=out_channel, kernel_size=(3, 3), padding=2)
        )
        self.adp = nn.AdaptiveMaxPool2d(output_size=1)
        self.classification = nn.Sequential(
            nn.Linear(in_features=out_channel, out_features=hidden_feature),
            nn.ReLU(),
            nn.Linear(in_features=hidden_feature, out_features=out_features),
            nn.Softmax(dim=1)
        )

    def forward(self, data):
        """
        前向传递函数
        """
        x, q = data
        x = self.laq(x, q)
        x.unsqueeze_(dim=1)  # 将输入数据变成(B, C, S, F) //B-batch size, C-channel size, S-sequence length, F-features
        x = self.cnn(x)
        x = self.adp(x)
        x = self.classification(x.squeeze())
        return x


if __name__ == '__main__':
    x = torch.rand((10, 20, 6))
    q = torch.rand((10, 3))
    g = torch.rand((3, 6))
    model = ICNN(g, 6, 1, 8, 16, 32, 64, 20)
    y = model((x, q))
    print(y.shape)
